import os
import sys
import json
import tempfile
import requests
from flask import Flask, request, jsonify
from flask_cors import CORS
import urllib.parse
import cv2
import numpy as np
import logging
import traceback
import threading

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger('wechat_extractor_api')

# 全局变量
reader = None
reader_lock = threading.Lock()
reader_loading = False

# 主应用
app = Flask(__name__)
CORS(app)  # 允许跨域请求
app.config['MAX_CONTENT_LENGTH'] = 50 * 1024 * 1024  # 限制最大上传大小为50MB
app.config['TIMEOUT'] = 120  # 设置请求超时时间为120秒

def load_easyocr_in_background():
    """在后台线程中加载EasyOCR模型"""
    global reader, reader_loading
    
    try:
        logger.info("开始在后台加载EasyOCR模型...")
        # 导入EasyOCR
        import easyocr
        # 初始化读取器
        reader = easyocr.Reader(['ch_sim', 'en'], gpu=True)
        logger.info("EasyOCR模型加载完成")
    except Exception as e:
        logger.error(f"加载EasyOCR模型失败: {e}")
    finally:
        reader_loading = False

def get_easyocr_reader():
    """获取或加载EasyOCR读取器"""
    global reader, reader_loading
    
    with reader_lock:
        if reader is not None:
            return reader
        
        if not reader_loading:
            reader_loading = True
            # 在后台线程加载模型
            threading.Thread(target=load_easyocr_in_background, daemon=True).start()
        
        return None

@app.route('/')
def home():
    return """
    <html>
        <head>
            <title>WeChat Dialog Extractor API</title>
            <style>
                body { font-family: Arial, sans-serif; line-height: 1.6; max-width: 800px; margin: 0 auto; padding: 20px; }
                code { background: #f4f4f4; padding: 2px 5px; border-radius: 3px; }
                pre { background: #f4f4f4; padding: 10px; border-radius: 5px; overflow-x: auto; }
                h1, h2 { border-bottom: 1px solid #eee; padding-bottom: 10px; }
                .example { margin: 20px 0; }
                .status { padding: 10px; margin: 10px 0; border-radius: 5px; }
                .ready { background-color: #dff0d8; }
                .loading { background-color: #fcf8e3; }
            </style>
        </head>
        <body>
            <h1>WeChat Dialog Extractor API</h1>
            <p>这是一个用于从微信聊天截图中提取对话内容的API服务。</p>
            
            <div id="status" class="status loading">
                <h3>模型状态:</h3>
                <p id="status-text">正在检查EasyOCR模型状态...</p>
            </div>
            <script>
                function checkStatus() {
                    fetch('/status')
                        .then(response => response.json())
                        .then(data => {
                            const statusDiv = document.getElementById('status');
                            const statusText = document.getElementById('status-text');
                            
                            if (data.ready) {
                                statusDiv.className = 'status ready';
                                statusText.textContent = '模型已加载完成，可以处理请求';
                            } else {
                                statusDiv.className = 'status loading';
                                statusText.textContent = '模型正在加载中，请稍后再试...';
                                setTimeout(checkStatus, 5000);
                            }
                        })
                        .catch(error => {
                            console.error('Error:', error);
                            document.getElementById('status-text').textContent = '检查状态失败，请刷新页面';
                        });
                }
                
                // 页面加载后检查状态
                window.onload = checkStatus;
            </script>
            
            <h2>API 用法</h2>
            <div class="example">
                <h3>请求示例</h3>
                <code>GET /extract?url=https://example.com/wechat_screenshot.jpg</code>
                
                <h3>参数</h3>
                <ul>
                    <li><code>url</code> - 图片的URL（必需）</li>
                </ul>
                
                <h3>响应格式</h3>
                <pre>{
  "success": true,
  "dialogues": [
    {
      "position": 0,
      "speaker": "用户",
      "content": "你好！"
    },
    {
      "position": 1,
      "speaker": "他人",
      "content": "你好，有什么可以帮助你的吗？"
    }
  ]
}</pre>
            </div>
            
            <h2>测试表单</h2>
            <form action="/extract" method="get">
                <label for="url">输入图片URL:</label><br>
                <input type="text" id="url" name="url" style="width: 100%; padding: 8px; margin: 10px 0;"><br>
                <input type="submit" value="提取对话" style="padding: 8px 15px; background: #4CAF50; color: white; border: none; cursor: pointer;">
            </form>
        </body>
    </html>
    """

@app.route('/status')
def check_status():
    """检查EasyOCR模型是否已加载"""
    global reader, reader_loading
    
    # 如果读取器未加载且未在加载中，触发加载
    if reader is None and not reader_loading:
        get_easyocr_reader()
    
    return jsonify({
        "ready": reader is not None,
        "loading": reader_loading
    })

@app.route('/extract')
def extract_from_url():
    """从URL提取对话内容"""
    image_url = request.args.get('url')
    
    if not image_url:
        return jsonify({
            'success': False,
            'error': '缺少URL参数'
        }), 400
    
    logger.info(f"收到提取请求，图片URL: {image_url}")
    
    # 检查EasyOCR是否已加载
    global reader
    if reader is None:
        # 尝试获取读取器
        reader = get_easyocr_reader()
        if reader is None:
            return jsonify({
                'success': False,
                'error': 'EasyOCR模型正在加载中，请稍后再试',
                'details': '首次使用需要下载模型，可能需要几分钟时间'
            }), 503  # Service Unavailable
    
    try:
        # 创建临时目录存储下载的图片
        with tempfile.TemporaryDirectory() as temp_dir:
            # 下载图片
            logger.info(f"开始下载图片...")
            temp_image_path = download_image(image_url, temp_dir)
            
            if not temp_image_path:
                return jsonify({
                    'success': False,
                    'error': '无法下载或保存图片'
                }), 400
            
            logger.info(f"图片已下载到: {temp_image_path}")
            
            # 处理图片并提取对话
            logger.info(f"开始处理图片并提取对话...")
            
            # 导入并使用处理函数
            from wechat_dialog_extractor import process_image_by_text_first
            
            # 替换全局reader
            import easyocr
            original_reader = easyocr.Reader
            easyocr.Reader = lambda *args, **kwargs: reader
            
            # 处理图像
            dialogues = process_image_by_text_first(temp_image_path)
            
            # 恢复原始reader
            easyocr.Reader = original_reader
            
            logger.info(f"成功提取 {len(dialogues)} 条对话")
            
            # 返回JSON结果
            return jsonify({
                'success': True,
                'dialogues': dialogues
            })
            
    except Exception as e:
        error_msg = str(e)
        stack_trace = traceback.format_exc()
        logger.error(f"处理出错: {error_msg}\n{stack_trace}")
        
        return jsonify({
            'success': False,
            'error': error_msg,
            'details': stack_trace
        }), 500

def download_image(url, save_dir):
    """下载图片并保存到临时目录"""
    try:
        # 发送GET请求获取图片
        logger.info(f"开始下载图片: {url}")
        response = requests.get(url, stream=True, timeout=30)  # 增加超时设置
        response.raise_for_status()
        
        # 从URL中提取图片文件名
        parsed_url = urllib.parse.urlparse(url)
        image_name = os.path.basename(parsed_url.path)
        
        # 如果文件名为空或无扩展名，使用默认文件名
        if not image_name or '.' not in image_name:
            image_name = 'image.jpg'
            
        # 确保文件名是有效的
        image_name = ''.join(c for c in image_name if c.isalnum() or c in '._-')
        if not image_name:
            image_name = 'image.jpg'
        
        # 保存图片
        image_path = os.path.join(save_dir, image_name)
        with open(image_path, 'wb') as file:
            for chunk in response.iter_content(chunk_size=8192):
                file.write(chunk)
        
        logger.info(f"图片成功下载到: {image_path}")
        return image_path
    
    except Exception as e:
        logger.error(f"下载图片失败: {str(e)}")
        return None

# 添加超时处理中间件
@app.after_request
def add_timeout_header(response):
    """增加超时相关的响应头"""
    response.headers['X-Accel-Buffering'] = 'no'  # 禁用Nginx缓冲
    return response

# 使用更新的方式处理首次请求
class InitModelMiddleware:
    def __init__(self, app):
        self.app = app
        self.initialized = False

    def __call__(self, environ, start_response):
        if not self.initialized:
            # 在第一次请求时触发模型加载
            get_easyocr_reader()
            self.initialized = True
        return self.app(environ, start_response)

# 添加中间件
app.wsgi_app = InitModelMiddleware(app.wsgi_app)

if __name__ == '__main__':
    # 获取命令行参数或使用默认值
    host = '0.0.0.0'  # 绑定到所有网络接口
    port = int(os.environ.get('PORT', 5000))
    debug = os.environ.get('DEBUG', 'False').lower() == 'true'
    
    # 在主线程中启动模型加载
    threading.Thread(target=get_easyocr_reader, daemon=True).start()
    
    logger.info(f"启动服务器 {host}:{port}, debug模式: {debug}")
    app.run(host=host, port=port, debug=debug)
